"""
中文期刊文章数据存储任务（Prefect Task）

本模块负责将一批中文期刊文章（以 pandas.DataFrame 形式传入）解析、清洗并批量写入 MySQL 关系型数据库。
整个流程在 **一个数据库事务** 内完成，任何一步失败都会自动回滚，保证数据一致性。

主要功能
1. 生成全局唯一 article_id。
2. 解析并保存文章基本信息（CnArticleList）。
3. 解析并保存文章类型字典（CnArticleTypeDict）及文章-类型关系（CnArticleTypeRelation）。
4. 解析并保存参考文献（CnArticleReferenceList）。
5. 解析并保存摘要（CnArticleAbstractInfo）。
6. 解析并保存关键词字典（CnArticleKeywordsList）及文章-关键词关系（CnArticleKeywordsRelation）。
7. 解析并保存基金信息（CnArticleGrantRelation）。

作者: <your-name>
更新: 2025-09-30
"""

import json
from collections import defaultdict
from datetime import datetime

import pandas as pd
from prefect import task, get_run_logger
from application.db import get_database_connection
from application.db.mysql_db.article_ch.CnArticleAbstractInfo import CnArticleAbstractInfo
from application.db.mysql_db.article_ch.CnArticleAuthorList import CnArticleAuthorList
from application.db.mysql_db.article_ch.CnArticleAuthorPlatformIdRelation import CnArticleAuthorPlatformIdRelation
from application.db.mysql_db.article_ch.CnArticleAuthorRelation import CnArticleAuthorRelation
from application.db.mysql_db.article_ch.CnArticleAuthorUnitList import CnArticleAuthorUnitList
from application.db.mysql_db.article_ch.CnArticleGrantRelation import CnArticleGrantRelation
from application.db.mysql_db.article_ch.CnArticleKeywordsList import CnArticleKeywordsList
from application.db.mysql_db.article_ch.CnArticleKeywordsRelation import CnArticleKeywordsRelation
from application.db.mysql_db.article_ch.CnArticleList import CnArticleList
from application.db.mysql_db.article_ch.CnArticleReferenceList import CnArticleReferenceList
from application.db.mysql_db.article_ch.CnArticleTypeDict import CnArticleTypeDict
from application.db.mysql_db.article_ch.CnArticleTypeRelation import CnArticleTypeRelation
from application.tasks.base_task import BaseTask
from application.utils.utils import generate_random_id


# --------------------------------------------------------------------------- #
# 主 Task 继承自 BaseTask，Prefect 会自动识别并注册
# --------------------------------------------------------------------------- #
class ArticleChDataStorageTask(BaseTask):
    """
    中文期刊数据落库任务
    """

    # ----------------------------------------------------------------------- #
    # 核心入口：run
    # ----------------------------------------------------------------------- #
    @get_database_connection("default").atomic()  # 整个函数在一个事务里执行
    def run(self, df: pd.DataFrame) -> pd.DataFrame:
        """
        主入口函数
        :param df: 待处理的文章 DataFrame，每行是一篇文章
        :return: 处理后的 DataFrame（原 df 上附加了各组件解析结果，方便调试）
        """
        logger = get_run_logger()
        logger.info("开始处理中文期刊数据，共 %s 条", len(df))

        # 1. 生成 article_id
        df["article_id"] = df.apply(lambda _: generate_random_id("cnlt"), axis=1)
        df['article_author'] = df.apply(self.article_author_load, axis=1)
        # 2. 逐列解析各组件
        df["article_type"] = df.apply(self.article_type_load, axis=1)
        df["article_info"] = df.apply(self.article_info_load, axis=1)
        df["article_reference"] = df.apply(self.article_reference_load, axis=1)
        df["article_abstract"] = df.apply(self.article_abstract_load, axis=1)
        df["article_keywords"] = df.apply(self.article_keywords_load, axis=1)
        df["article_grant"] = df.apply(self.article_grant_load, axis=1)

        # 3. 将嵌套列拆成独立 DataFrame，方便批量插入
        article_info = (
            df["article_info"]
            .apply(pd.Series)
            .dropna(how="all")
            .dropna(how="all", axis=1)
        )
        article_reference = (
            df["article_reference"]
            .apply(pd.Series)
            .dropna(how="all")
            .dropna(how="all", axis=1)
        )
        article_abstract = (
            df["article_abstract"]
            .apply(pd.Series)
            .dropna(how="all")
            .dropna(how="all", axis=1)
        )
        article_keywords = (
            df["article_keywords"]
            .explode()
            .apply(pd.Series)
            .dropna(how="all")
            .dropna(how="all", axis=1)
        )
        article_grant = (
            df["article_grant"]
            .explode()
            .apply(pd.Series)
            .dropna(how="all")
            .dropna(how="all", axis=1)
        )
        article_type = (
            df["article_type"]
            .apply(pd.Series)
            .dropna(how="all")
            .dropna(how="all", axis=1)
        )
        # 从原始DataFrame中提取“article_author”列，并进行初步清洗和展开
        article_author = (
            df["article_author"]
            .explode()  # 将列表类型的单元格元素展开为多行
            .apply(pd.Series)  # 将展开后的每个元素转化为独立的列（Series -> DataFrame）
            .dropna(how="all")  # 删除整行全为空值的数据
            .dropna(how="all", axis=1)  # 删除整列全为空值的数据
        )

        # 修复逻辑：确保在数据存在的情况下再进行后续处理，避免空数据导致错误
        if not article_author.empty:
            # 对每一行作者数据执行 transformation_author_super 函数进行预处理
            # 之后删除空值行，再将结果展开为多列，并继续调用 get_author 函数进行进一步提取
            article_author = (
                article_author
                .apply(self.article_transformation_author_super, axis=1)  # 应用作者信息的预处理转换函数
                .dropna()  # 删除转换后为空的记录
                .apply(pd.Series)  # 将嵌套结构再次展开成列
                .apply(self.article_get_author, axis=1)  # 提取作者核心信息（如姓名、单位等）
            )

            # 修复：防止调用 explode(0) 导致 KeyError: 0 错误
            if not article_author.empty and len(article_author) > 0:
                # 尝试展开作者列中可能存在的列表结构（如多作者情况）
                exploded = article_author.explode()
                if not exploded.empty:
                    # 若展开结果非空，则去除其中的空值
                    article_author = exploded.dropna()
                else:
                    # 若展开结果为空，则返回空的 Series（保持数据结构一致）
                    article_author = pd.Series([], dtype=object)
            else:
                # 若 article_author 为空或无有效行，返回空的 Series
                article_author = pd.Series([], dtype=object)
        else:
            # 若初始 article_author 为空，则直接返回空的 Series
            article_author = pd.Series([], dtype=object)

        # 4. 批量写入数据库
        # 4.1 文章主表
        if not article_info.empty:
            CnArticleList.insert_many(
                self.clean_dataframe(article_info).to_dict(orient="records")
            ).execute()

        # 4.2 文章类型字典（仅新增）
        if not article_type.empty:
            CnArticleTypeDict.insert_many(
                self.clean_dataframe(article_type[article_type["exist"] == False][
                                         ["type_code", "type_name", "type_name_eng", "type_creator"]
                                     ]).to_dict(orient="records")
            ).execute()

        # 4.3 文章-类型关系
        if not article_type.empty:
            CnArticleTypeRelation.insert_many(
                self.clean_dataframe(
                    article_type[["article_id", "type_code"]]
                ).to_dict(orient="records")
            ).execute()

        # 4.4 参考文献
        if not article_reference.empty:
            CnArticleReferenceList.insert_many(
                self.clean_dataframe(article_reference).to_dict(orient="records")
            ).execute()

        # 4.5 摘要
        if not article_abstract.empty:
            CnArticleAbstractInfo.insert_many(
                self.clean_dataframe(article_abstract).to_dict(orient="records")
            ).execute()

        # 4.6 关键词字典（仅新增）
        if not article_keywords.empty:
            new_keywords = article_keywords.loc[
                article_keywords["exist"] == False, ["keyword_id", "keyword_chs", "keyword_eng"]]
            if not new_keywords.empty:
                CnArticleKeywordsList.insert_many(
                    self.clean_dataframe(new_keywords).to_dict(orient="records")
                ).execute()

        # 4.7 文章-关键词关系
        if not article_keywords.empty:
            CnArticleKeywordsRelation.insert_many(
                self.clean_dataframe(
                    article_keywords[["article_id", "keyword_id"]]
                ).to_dict(orient="records")
            ).execute()

        # 4.8 基金关系
        if not article_grant.empty:
            CnArticleGrantRelation.insert_many(
                self.clean_dataframe(article_grant).to_dict(orient="records")
            ).execute()
        # 4.9 作者表
        if not article_author.empty:
            CnArticleAuthorRelation.insert_many(article_author.to_list()).execute()
        logger.info(f"中文期刊数据落库完成,{len(df)} 行数据处理完毕")
        return df

    # ----------------------------------------------------------------------- #
    # 工具函数：清洗 DataFrame，把 NaN 统一转成 None，避免 MySQL 驱动报错
    # ----------------------------------------------------------------------- #
    def clean_dataframe(self, dataframe: pd.DataFrame) -> pd.DataFrame:
        """
        将 DataFrame 中的 NaN 值替换为 None，避免 MySQL 插入错误
        """
        return dataframe.where(pd.notnull(dataframe), None)

    # ----------------------------------------------------------------------- #
    # 以下依次为各组件的解析函数，返回 dict 或 list[dict]，供 apply 调用
    # ----------------------------------------------------------------------- #
    def article_info_load(self, item: pd.Series) -> dict:
        """
        解析文章基本信息
        """
        item = item.dropna()
        article_info = item.get("article_info") or {}
        article_id = item["article_id"]
        article_name = article_info.get("article_name") or {}
        # 其余字段...
        return {
            "article_id": article_id,
            "article_name": article_name,
            "article_doi": article_info.get("article_doi"),
            "journal_code": article_info.get("journal_info", {}).get("journal_code"),
            "article_year": article_info.get("article_year"),
            "article_language": article_info.get("article_language"),
            "original_file": article_info.get("original_file"),
            "file_size": article_info.get("file_size"),
            "upload_time": article_info.get("upload_time"),
            "published_time": article_info.get("published_time"),
            "published_time_precision": article_info.get("published_time_precision"),
            "article_time": article_info.get("article_time"),
            "published_info": article_info.get("published_info"),
            "clc_codes": article_info.get("clc_codes"),
            "platform_id": article_info.get("platform_id"),
            "article_creator": article_info.get("platform_source"),
            "revise_time": datetime.now(),
        }

    def article_reference_load(self, item: pd.Series) -> dict:
        """
        解析参考文献
        """
        item = item.dropna()
        return {
            "article_id": item["article_id"],
            "reference_info": item.get("article_reference") or {},
        }

    def article_abstract_load(self, item: pd.Series) -> dict:
        """
        解析摘要
        """
        return {
            "article_id": item["article_id"],
            "chs_text": (item.get("article_abstract") or {}).get("chs_info") or [],
            "eng_text": (item.get("article_abstract") or {}).get("eng_info") or [],
        }

    def article_keywords_load(self, item: pd.Series) -> list[dict]:
        """
        解析关键词，同时与已有关键词字典比对，生成 keyword_id 并标记是否已存在
        """
        item = item.dropna()
        article_id = item["article_id"]
        article_keywords = item.get("article_keywords") or []
        if not article_keywords:
            return []

        # 已有关键词映射
        keywords_set = {kw["keyword_chs"] for kw in article_keywords}
        exist_kw = {
            kw.keyword_chs: kw.keyword_id
            for kw in CnArticleKeywordsList.select(
                CnArticleKeywordsList.keyword_id,
                CnArticleKeywordsList.keyword_chs,
            ).where(CnArticleKeywordsList.keyword_chs.in_(keywords_set))
        }

        # 补全字段
        for kw in article_keywords:
            chs = kw["keyword_chs"]
            if chs in exist_kw:
                kw["keyword_id"] = exist_kw[chs]
                kw["exist"] = True
            else:
                kw["keyword_id"] = generate_random_id("cnky")
                kw["exist"] = False
            kw["article_id"] = article_id
        return article_keywords

    def article_grant_load(self, item: pd.Series) -> list[dict]:
        """
        解析基金信息
        """
        item = item.dropna()
        article_id = item["article_id"]
        article_grant = item.get("article_grant")
        if not isinstance(article_grant, list):
            return []

        result = []
        for idx, grant in enumerate(article_grant):
            result.append(
                {
                    "article_id": article_id,
                    "grant_name": grant.get("grant_name"),
                    "grant_order": grant.get("grant_order") or idx + 1,
                    "platform_id": grant.get("platform_id"),
                    "standard_grant": None,  # 预留字段
                }
            )
        return result

    def article_type_load(self, item: pd.Series) -> dict:
        """
        解析文章类型，同时与已有类型字典比对
        """
        item = item.dropna()
        article_id = item["article_id"]
        article_info = item.get("article_info") or {}
        type_name = article_info.get("type_name") or {}
        type_cn = type_name.get("chi", "")
        type_en = type_name.get("eng", "")

        # 查询是否已存在
        exist = (
            CnArticleTypeDict.select(
                CnArticleTypeDict.type_code,
                CnArticleTypeDict.type_name,
                CnArticleTypeDict.type_name_eng,
            )
            .where(
                (CnArticleTypeDict.type_name == type_cn)
                | (CnArticleTypeDict.type_name_eng == type_en)
            )
            .first()
        )

        if exist:
            return {
                "article_id": article_id,
                "type_code": exist.type_code,
                "type_name": exist.type_name,
                "type_name_eng": exist.type_name_eng,
                "type_creator": article_info.get("platform_source"),
                "exist": True,
            }
        else:
            new_code = generate_random_id("cnat")
            return {
                "article_id": article_id,
                "type_code": new_code,
                "type_name": type_cn,
                "type_name_eng": type_en,
                "type_creator": article_info.get("platform_source"),
                "exist": False,
            }

    def article_author_load(self, item):
        item = self.clean_dataframe(item)
        platform_code = item.get("article_info").get('platform_source')
        article_id = item.get('article_id')
        author_list = item.get("article_author", []) or []
        result = []
        for index, author in enumerate(author_list):
            insert = author | {"article_id": article_id, "author_index": index, "platform_code": platform_code}
            result.append(insert)
        return result

    def article_get_author(self, item):
        item = self.clean_dataframe(item)
        result_list = []
        unit_info = item.get('unit_info') or []
        for index, i in enumerate(unit_info):
            result_list.append({
                'article_id': item.get('article_id'),
                'author_id': item.get('author_id'),
                'unit_id': i.get('unit_id'),
                "author_order": item.get('author_index'),
                'unit_order': index,
                'corresponding_author': item.get('corresponding_author'),
                "author_email": item.get('author_email'),
                "is_display": item.get('is_display', 1),
            })
        return result_list

    def article_transformation_author_super(self, author):
        author = self.clean_dataframe(author)
        platform_id = author.get('platform_info', {}).get('platform_id')

        unit_info = author.get('unit_info') or []

        unit_info_names = [i.get('unit_name') for i in unit_info]
        platform_code = author.get('platform_code')
        unit_info_name_dict = {}
        cn_a_u = CnArticleAuthorUnitList.select(CnArticleAuthorUnitList.unit_id,
                                                CnArticleAuthorUnitList.unit_name).where(
            CnArticleAuthorUnitList.unit_name.in_(unit_info_names)).dicts()
        for i in cn_a_u:
            unit_info_name_dict[i.get('unit_name')] = i.get('unit_id')

        for unit in unit_info:

            if unit.get('unit_name') in unit_info_name_dict:
                unit['unit_id'] = unit_info_name_dict[unit.get('unit_name')]
            else:
                unit['unit_id'] = generate_random_id('JG')
                unit_info_name_dict[unit.get('unit_name')] = unit['unit_id']
                insert_unit = {
                    'unit_id': unit['unit_id'],
                    'unit_name': unit['unit_name']
                }
                if unit.get('platform_id'):
                    insert_unit['platform_id'] = {platform_code: unit.get('platform_id')}

                CnArticleAuthorUnitList.insert(insert_unit).execute()

        if platform_id:
            exists = CnArticleAuthorPlatformIdRelation.select(
                CnArticleAuthorPlatformIdRelation.author_id
            ).where(
                CnArticleAuthorPlatformIdRelation.platform_id == platform_id
            ).first()
            if exists:
                author['author_id'] = exists.author_id

        if not author.get('author_id'):
            cn_a_l = CnArticleAuthorList.select(CnArticleAuthorList.author_id).where(
                CnArticleAuthorList.author_name == author.get('author_name')
            ).dicts()
            author_ids = [i.get('author_id') for i in cn_a_l]
            cn_a_r = CnArticleAuthorRelation.select(CnArticleAuthorRelation.author_id,
                                                    CnArticleAuthorRelation.unit_id).where(
                CnArticleAuthorRelation.author_id.in_(author_ids)
            ).dicts()
            for i in cn_a_r:
                if i.get('unit_id') in unit_info_name_dict.values():
                    author['author_id'] = i.get('author_id')
                    break

        if not author.get('author_id'):
            author['author_id'] = generate_random_id('cn')
            CnArticleAuthorList.insert({
                'author_id': author.get('author_id'),
                'author_name': author.get('author_name'),
                'author_name_eng': author.get('author_name_eng'),
                'last_name': author.get('last_name'),
                'fore_name': author.get('fore_name'),
                'author_type': author.get('author_type'),
                'author_desc': ''
            }).execute()
            if platform_id:
                CnArticleAuthorPlatformIdRelation.insert({
                    'author_id': author.get('author_id'),
                    'platform_id': platform_id,
                    'platform_code': author.get('platform_info').get('platform_code')
                }).execute()
        return author
